import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import uuid

# 初始化参数
nx, ny = 100, 100  # 网格尺寸
dt = 0.01          # 时间步长
t_max = 10.0       # 最大时间
rho = 1.0          # 流体密度
nu = 0.1           # 运动粘度

# PID 参数
K_p = 1.0          # 比例增益
K_i = 0.1          # 积分增益
K_d = 0.01         # 微分增益

# 初始化速度场和目标速度场
u = np.zeros((nx, ny, 2))  # 2D 速度场
u_target = np.ones((nx, ny, 2)) * 0.5  # 目标速度场

# 初始化 PID 控制器的积分和微分项
e_integral = np.zeros_like(u)
e_prev = np.zeros_like(u)

# 用于可视化的网格
X, Y = np.meshgrid(np.arange(nx), np.arange(ny))
frames = []  # 存储每帧的矢量图数据
plot_interval = 10  # 每 10 步保存一帧

# 定义 Navier-Stokes 方程的右端项
def navier_stokes_rhs(t, u_flat):
    u = u_flat.reshape((nx, ny, 2))
    
    # 计算误差
    e = u_target - u
    
    # 更新 PID 控制器的积分和微分项
    global e_integral, e_prev
    e_integral += e * dt
    e_derivative = (e - e_prev) / dt
    e_prev = e
    
    # 计算 PID 控制器的输出
    u_pid = K_p * e + K_i * e_integral + K_d * e_derivative
    
    # 分离速度分量
    u_x = u[..., 0]  # x 方向速度分量
    u_y = u[..., 1]  # y 方向速度分量
    
    # 计算对流项 (u · ∇)u
    grad_ux = np.gradient(u_x, axis=(0, 1))  # ∂u_x/∂x, ∂u_x/∂y
    grad_uy = np.gradient(u_y, axis=(0, 1))  # ∂u_y/∂x, ∂u_y/∂y
    
    # 对流项：u_x * (∂u_x/∂x, ∂u_y/∂x) + u_y * (∂u_x/∂y, ∂u_y/∂y)
    convection_x = u_x * grad_ux[0] + u_y * grad_ux[1]
    convection_y = u_x * grad_uy[0] + u_y * grad_uy[1]
    convection = np.stack([convection_x, convection_y], axis=-1)
    
    # 计算压力梯度 (假设压力场 p 为零)
    p = np.zeros((nx, ny))
    grad_p = np.stack(np.gradient(p, axis=(0, 1)), axis=-1)
    
    # 计算粘性项 ν ∇²u
    laplacian_ux = np.zeros_like(u_x)
    laplacian_uy = np.zeros_like(u_y)
    
    # 对 x 分量和 y 分量分别计算拉普拉斯算子
    grad_ux_x = np.gradient(grad_ux[0], axis=0)  # ∂²u_x/∂x²
    grad_ux_y = np.gradient(grad_ux[1], axis=1)  # ∂²u_x/∂y²
    laplacian_ux = grad_ux_x + grad_ux_y
    
    grad_uy_x = np.gradient(grad_uy[0], axis=0)  # ∂²u_y/∂x²
    grad_uy_y = np.gradient(grad_uy[1], axis=1)  # ∂²u_y/∂y²
    laplacian_uy = grad_uy_x + grad_uy_y
    
    laplacian_u = np.stack([laplacian_ux, laplacian_uy], axis=-1)
    
    # 计算 Navier-Stokes 方程的右端项
    du_dt = -convection - 1 / rho * grad_p + nu * laplacian_u + u_pid
    
    # 确保 du_dt 的形状与 u 的形状一致
    assert du_dt.shape == u.shape, f"du_dt shape {du_dt.shape} does not match u shape {u.shape}"
    
    # 展平 du_dt
    return du_dt.flatten()

# 生成 Graphviz 可视化流程图
def generate_graphviz_flowchart():
    dot = """
    digraph NavierStokes {
        rankdir=TB;
        node [shape=box, style=filled, fillcolor=lightblue];
        edge [color=navy];

        // 输入
        u [label="Velocity Field\\n(u, shape: 100x100x2)"];
        u_target [label="Target Velocity\\n(u_target, shape: 100x100x2)"];

        // PID 控制器
        error [label="Error\\n(e = u_target - u)"];
        integral [label="Integral Term\\n(e_integral += e * dt)"];
        derivative [label="Derivative Term\\n(e_derivative = (e - e_prev)/dt)"];
        pid [label="PID Output\\n(u_pid = K_p*e + K_i*e_integral + K_d*e_derivative)"];

        // 对流项
        split [label="Split Components\\n(u_x, u_y)"];
        grad_ux [label="Gradient of u_x\\n(∂u_x/∂x, ∂u_x/∂y)"];
        grad_uy [label="Gradient of u_y\\n(∂u_y/∂x, ∂u_y/∂y)"];
        convection [label="Convection Term\\n(u_x * grad_ux + u_y * grad_uy)"];

        // 压力梯度
        pressure [label="Pressure Field\\n(p, shape: 100x100)"];
        grad_p [label="Pressure Gradient\\n(∇p, shape: 100x100x2)"];

        // 粘性项
        laplacian_ux [label="Laplacian of u_x\\n(∂²u_x/∂x² + ∂²u_x/∂y²)"];
        laplacian_uy [label="Laplacian of u_y\\n(∂²u_y/∂x² + ∂²u_y/∂y²)"];
        laplacian [label="Viscous Term\\n(ν * ∇²u, shape: 100x100x2)"];

        // 最终输出
        du_dt [label="Time Derivative\\n(du_dt = -convection - 1/ρ * grad_p + ν * laplacian + u_pid)"];
        output [label="Flattened Output\\n(du_dt.flatten())"];

        // 数据流
        u -> error;
        u_target -> error;
        error -> integral;
        error -> derivative;
        error -> pid;
        integral -> pid;
        derivative -> pid;

        u -> split;
        split -> grad_ux;
        split -> grad_uy;
        grad_ux -> convection;
        grad_uy -> convection;

        pressure -> grad_p;

        grad_ux -> laplacian_ux;
        grad_uy -> laplacian_uy;
        laplacian_ux -> laplacian;
        laplacian_uy -> laplacian;

        convection -> du_dt;
        grad_p -> du_dt;
        laplacian -> du_dt;
        pid -> du_dt;
        du_dt -> output;
    }
    """
    return dot

# 主循环与可视化
fig, ax = plt.subplots(figsize=(8, 8))
t = 0.0
step = 0

def update_quiver(frame):
    ax.clear()
    u_x = frame[..., 0]
    u_y = frame[..., 1]
    # 为了减少计算量，使用稀疏网格绘制矢量
    stride = 5
    ax.quiver(X[::stride, ::stride], Y[::stride, ::stride], 
              u_x[::stride, ::stride], u_y[::stride, ::stride], 
              scale=20, color='blue')
    ax.set_title(f"Velocity Field at t = {t:.2f}")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_xlim(0, nx)
    ax.set_ylim(0, ny)
    ax.grid(True)

while t < t_max:
    # 更新流体动力学
    sol_u = solve_ivp(navier_stokes_rhs, [t, t + dt], u.flatten(), method='RK45')
    u = sol_u.y[:, -1].reshape((nx, ny, 2))
    
    # 每 plot_interval 步保存一帧
    if step % plot_interval == 0:
        frames.append(u.copy())
    
    # 更新时间和步数
    t += dt
    step += 1

# 生成动画
ani = FuncAnimation(fig, update_quiver, frames=frames, interval=100)
ani.save("velocity_field_evolution.gif", writer=PillowWriter(fps=10))
plt.close()

# 保存 Graphviz 流程图
graphviz_content = generate_graphviz_flowchart()
with open("navier_stokes_flowchart.dot", "w") as f:
    f.write(graphviz_content)

print("Navier-Stokes simulation completed.")
print("Graphviz flowchart saved as 'navier_stokes_flowchart.dot'.")
print("Velocity field animation saved as 'velocity_field_evolution.gif'.")